﻿// Copyright 2009 Apoc @ ApocDev.com | Apoc @ MMOwned.com | Apoc @ GameDeception.net
using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Runtime.InteropServices;
using System.Xml.Linq;

namespace RoMFramework
{
    /// <summary>
    /// Credits to Dominik, Patrick, Bobbysing, and whoever else I forgot, for most of the ripped code here!
    /// </summary>
    public class FindPattern
    {
        /// <summary>
        /// Pointers to dynamic addresses or moar pointers
        /// </summary>
        public readonly Dictionary<string, uint> PointersDictionary = new Dictionary<string, uint>();
        /// <summary>
        /// Static Addresses
        /// </summary>
        public readonly Dictionary<string, uint> AddressesDictionary = new Dictionary<string, uint>();
        protected byte[] data;

        /// <summary>
        /// Creates a new instance of the <see cref="FindPattern"/> class. This class will read from a specified patterns XML file
        /// and search out those patterns in the specified process's memory.
        /// </summary>
        /// <param name="patternFile">The full path to the pattern XML file.</param>
        /// <param name="processHandle">An open process handle to the process to read memory from.</param>
        /// <param name="startAddr">The 'base' address of the process (or module)</param>
        /// <param name="endAddr">The 'end' of the process (or module). Eg; where to stop reading memory from.</param>
        public FindPattern(string patternFile, IntPtr processHandle, uint startAddr, uint endAddr)
        {
            // Get a temporary set of data to work with. :)
            byte[] data = ReadBytes(processHandle, (IntPtr)startAddr, (int)(endAddr - startAddr));
            LoadFile(XElement.Load(patternFile), data, startAddr);
        }

        /// <summary>
        /// Creates a new instance of the <see cref="FindPattern"/> class. This class will read from a specified patterns XML file
        /// and search out those patterns in the specified process's memory.
        /// </summary>
        /// <param name="patternFile">The full path to the pattern XML file.</param>
        /// <param name="startAddr">The 'base' address of the process (or module)</param>
        /// <param name="d">The bytes to search through</param>
        public FindPattern(string patternFile, uint startAddr, ref byte[] d)
        {
            data = d;
            LoadFile(XElement.Load(patternFile), data, startAddr);
        }

        /// <summary>
        /// Retrieves an address from the found patterns stash.
        /// </summary>
        /// <param name="name">The name of the pattern, as per the XML file provided in the constructor of this class instance.</param>
        /// <returns></returns>
        public uint this[string name] { get { return Get(name); } }

        [DllImport("kernel32.dll", SetLastError = true)]
        protected static extern bool ReadProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress, [Out] byte[] lpBuffer,
                                                     int dwSize, out int lpNumberOfBytesRead);

        protected static byte[] ReadBytes(IntPtr processHandle, IntPtr address, int count)
        {
            var ret = new byte[count];
            int numRead;
            if (ReadProcessMemory(processHandle, address, ret, count, out numRead) && numRead == count)
            {
                return ret;
            }
            return null;
        }

        /// <summary>
        /// Retrieves an address from the found patterns stash.
        /// </summary>
        /// <param name="name">The name of the pattern, as per the XML file provided in the constructor of this class instance.</param>
        /// <returns></returns>
        public uint Get(string name)
        {
            return PointersDictionary[name];
        }

        protected void LoadFile(XContainer file, byte[] data, uint start)
        {
            // Grab all the <Pattern /> elements from the XML.
            IEnumerable<XElement> pats = from p in file.Descendants("Pattern")
                               select p;

            // Each Pattern element needs to be handled seperately.
            // The enumeration we're goinv over, is in document order, so attributes such as 'start'
            // should work perfectly fine.
            foreach (XElement pat in pats)
            {
                uint tmpStart = 0;

                string usage = pat.Attribute("usage").Value;
                string name = pat.Attribute("desc").Value;
                string mask = pat.Attribute("mask").Value;
                byte[] patternBytes = GetBytesFromPattern(pat.Attribute("pattern").Value);

                // Make sure we're not getting some sort of screwy XML data.
                if (mask.Length != patternBytes.Length)
                    throw new Exception("Pattern and mask lengths do not match!");

                // If we run into a 'start' attribute, we need to remember that we're working from a 0
                // based 'memory pool'. So we just remove the 'start' from the address we found earlier.
                if (pat.Attribute("start") != null)
                {
                    tmpStart = Get(pat.Attribute("start").Value) - start + 1;
                }

                // Actually search for the pattern match...
                uint found = Find(data, mask, patternBytes, tmpStart);

                if (found == 0)
                    throw new Exception("Couldn't find a pattern!");
                if (pat.Elements() == null)
                {
                }

                // Handle specific child elements for the pattern.
                // <Lea> <Rel> <Add> <Sub> etc
                foreach (XElement e in pat.Elements())
                {
                    switch (e.Name.LocalName)
                    {
                        case "Lea":
                            found = BitConverter.ToUInt32(data, (int)found);
                            break;
                        case "Rel":
                            int instructionSize = int.Parse(e.Attribute("size").Value, NumberStyles.HexNumber);
                            int operandOffset = int.Parse(e.Attribute("offset").Value, NumberStyles.HexNumber);
                            found = (uint)(BitConverter.ToUInt32(data, (int)found) + found + instructionSize - operandOffset);
                            break;
                        case "Add":
                            found += uint.Parse(e.Attribute("value").Value, NumberStyles.HexNumber);
                            break;
                        case "Sub":
                            found -= uint.Parse(e.Attribute("value").Value, NumberStyles.HexNumber);
                            break;
                    }
                }

                //Adds the found pattern to addresses or pointers
                switch (usage)
                {
                    case "Pointer":
                        PointersDictionary.Add(name, found + start);
                        break;
                    case "Address":
                        AddressesDictionary.Add(name, found + start);
                        break;
                        //If the user didn't specify anything, or an invalid value at it to pointers by default
                    default:
                        PointersDictionary.Add(name, found + start);
                        break;

                }

            }
        }

        protected static byte[] GetBytesFromPattern(string pattern)
        {
            // Because I'm lazy, and this just makes life easier.
            string[] split = pattern.Split(new[] { '\\', 'x' }, StringSplitOptions.RemoveEmptyEntries);
            var ret = new byte[split.Length];
            for (int i = 0; i < split.Length; i++)
            {
                ret[i] = byte.Parse(split[i], NumberStyles.HexNumber);
            }
            return ret;
        }

        protected static uint Find(byte[] data, string mask, byte[] byteMask, uint start)
        {
            // There *has* to be a better way to do this stuff,
            // but for now, we'll deal with it.
            if (data != null)
            {
                for (uint i = start; i < data.Length; i++)
                {
                    if (DataCompare(data, (int)i, byteMask, mask))
                        return i;
                }
            }

            return 0;
        }

        protected static bool DataCompare(byte[] data, int offset, byte[] byteMask, string mask)
        {
            // Only check for 'x' mismatches. As we'll assume anything else is a wildcard.
            for (int i = 0; i < mask.Length; i++)
            {
                if (mask[i] == 'x' && byteMask[i] != data[i + offset])
                {
                    return false;
                }
            }
            return true;
        }
    }
}